"""
    IMPORTING LIBS
"""
import dgl

import numpy as np
import os
import socket
import time
import random
import glob
import argparse, json
import pickle

import torch
import torch.nn as nn
import torch.nn.functional as F

from dgl.dataloading import GraphDataLoader
from ogb.graphproppred import DglGraphPropPredDataset, Evaluator
import torch.optim as optim
from torch.utils.data import DataLoader

from tensorboardX import SummaryWriter
from tqdm import tqdm

# 这个类的作用就是创建一个字典，让它的键值对可以通过点号（.）来访问，就像访问对象的属性一样。
class DotDict(dict):
    # __init__(self, **kwds) 方法是类的构造函数，用于初始化对象。它接受任意数量的关键字参数 kwds。
    def __init__(self, **kwds):
        # self.update(kwds) 将传入的关键字参数更新到当前对象中。
        self.update(kwds)
        # self.__dict__ = self 将对象本身设为其自己的 __dict__ 属性，这样就可以让对象的属性像普通属性一样被访问。
        self.__dict__ = self






"""
    IMPORTING CUSTOM MODULES/METHODS
"""
from nets.superpixels_graph_classification.load_net import gnn_model # import all GNNS
from data.data import LoadData # import dataset




"""
    GPU Setup
"""
def gpu_setup(use_gpu, gpu_id):
    # 这两行代码的作用是告诉 CUDA 可见的 GPU 设备是哪一个，以便于程序在需要时正确选择并使用 GPU。
    # os.environ 是一个字典，用于获取和设置环境变量。
    # 将环境变量 CUDA_DEVICE_ORDER 设置为 "PCI_BUS_ID"，表示在选择可用的 GPU 设备时按照 PCI 总线 ID 的顺序进行选择。
    # 将环境变量 CUDA_VISIBLE_DEVICES 设置为函数参数 gpu_id 的字符串表示形式，用于指定可见的 GPU 设备。
    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)  

    # 这一行检查当前环境是否支持 CUDA（即是否有可用的 GPU）并且用户是否希望使用 GPU。
    if torch.cuda.is_available() and use_gpu:
        print('cuda available with GPU:',torch.cuda.get_device_name(0))
        device = torch.device("cuda")
    else:
        print('cuda not available')
        device = torch.device("cpu")
    return device










"""
    VIEWING MODEL CONFIG AND PARAMS
    查看模型配置和参数
"""
def view_model_param(MODEL_NAME, net_params):
    # 在这一行中，我们调用了一个叫做 gnn_model 的函数，并传入了 MODEL_NAME 和 net_params 作为参数，以此来创建一个模型对象 model。
    model = gnn_model(MODEL_NAME, net_params)
    # 我们初始化了一个变量 total_param，用于计算模型的总参数数量。
    total_param = 0
    # 这一行简单地打印出一个字符串，用于提醒用户接下来会显示模型的详细信息。
    print("model details:\n")
    #print(model)
    # 这是一个 for 循环，用于迭代模型的所有参数。
    for param in model.parameters():
        # print(param.data.size())
        # 在每次迭代中，我们将参数的大小转换为一个列表，并使用 NumPy 的 np.prod() 函数计算列表中所有元素的乘积，即参数的数量，并将其加到 total_param 上。
        total_param += np.prod(list(param.data.size()))
    # 最后，我们打印出模型名称和总参数数量。
    print('model/Total parameters:', MODEL_NAME, total_param)
    return total_param


"""
    TRAINING CODE
    创建要保存的检查点的文件名称
"""
# 这一行定义了一个名为 create_filename 的函数，它有三个参数：save_dir 表示保存目录，isbest 表示是否是最佳模型，num_epochs 表示模型的训练轮数。
def create_filename(save_dir, isbest=False, num_epochs=-1):
    """
    Args:
        args        :  the arguments parsed in the parser
        isbest      :  whether the saved model is the best-performing one
        num_epochs  :  epoch number of the model (when isbest=False)
    """
    # os.path.join() 函数用于拼接路径，这里拼接了 save_dir 和模型的基本名称 'MNIST_GCN_0.999'，然后将结果赋给 filename。
    filename = os.path.join(save_dir, 'MNIST_GCN_0.999')
    # os.makedirs() 函数用于创建目录，exist_ok=True 表示如果目录已经存在，则不会抛出异常。这行代码的作用是创建保存模型的目录。
    os.makedirs(filename, exist_ok=True)

    # 这段代码根据 isbest 和 num_epochs 参数来修改 filename。
    # 如果 isbest 为 True，则在 filename 后面添加 "best" 字符串；如果 num_epochs 大于 0，则在 filename 后面添加 num_epochs 的字符串表示。
    if isbest:
        filename = os.path.join(filename, "best")
    elif num_epochs > 0:
        filename = os.path.join(filename, str(num_epochs))

    # 最后，函数返回拼接完成的文件名，加上了后缀 "_undirected.pth.tar"。
    return filename + "_undirected.pth.tar"

# 这一行定义了一个名为 save_checkpoint 的函数，它接受多个参数：
# model 是要保存的模型，optimizer 是模型的优化器，ckptdir 是保存模型的目录，model_name 是模型的名称，
# num_epochs 是训练轮数，默认为 -1，isbest 表示是否为最佳模型，默认为 False，cg_dict 是计算图的字典，默认为 None。
def save_checkpoint(model, optimizer, ckptdir, model_name, num_epochs=-1, isbest=False, cg_dict=None):
    """Save pytorch model checkpoint.

    Args:
        - model         : The PyTorch model to save.
        - optimizer     : The optimizer used to train the model.
        - args          : A dict of meta-data about the model.
        - num_epochs    : Number of training epochs.
        - isbest        : True if the model has the highest accuracy so far.
        - cg_dict       : A dictionary of the sampled computation graphs.
    """
    # 调用 create_filename() 函数创建要保存的文件名，根据 ckptdir、isbest 和 num_epochs 参数确定是否添加额外的标识。
    filename = create_filename(ckptdir, isbest, num_epochs=num_epochs)
    # 使用 PyTorch 的 torch.save() 函数将模型、优化器和其他必要的信息保存到文件中。
    # 保存的内容是一个字典，包括当前的训练轮数 num_epochs、模型的名称 model_name、优化器 optimizer 的状态字典、模型的状态字典 model.state_dict() 和优化器的状态字典 optimizer.state_dict()。
    torch.save(
        {
            "epoch": num_epochs,
            "model_type": model_name,
            "optimizer": optimizer,
            "model_state": model.state_dict(),
            "optimizer_state": optimizer.state_dict(),
        },
        filename,
    )

# 这是一个评估模型性能的函数。
# 它接受模型 model、设备 device、数据加载器 loader、评估器 evaluator 以及两个额外的评估器 Acc_evaluator 和 F1_evaluator。
def eval(model, device, loader, evaluator, Acc_evaluator, F1_evaluator):
    # 将模型设置为评估模式，这样模型中的 dropout 等层将被禁用。
    model.eval()
    # 初始化两个列表，用于存储真实标签和模型预测标签。
    y_true = []
    y_pred = []

    # 使用 loader 遍历数据集的每个批次。
    for step, batch in enumerate(loader):
        # 将当前批次的图数据移动到指定的设备上。
        batch_graphs = batch[0].to(device)
        # 如果当前批次只包含一个图，则跳过，因为某些模型可能需要多个图进行处理。
        if batch_graphs.ndata['feat'].size(0) == 1:
            pass
        else:
            # 使用 torch.no_grad() 上下文管理器，禁用梯度计算。
            with torch.no_grad():
                # 将当前批次的节点特征和边特征转换为浮点型，并将它们移动到指定的设备上。
                batch_x = batch_graphs.ndata['feat'].float().to(device)  # num x feat
                batch_e = batch_graphs.edata['feat'].to(device)
                # 通过模型前向传播，得到预测结果 pred 和嵌入向量 embed。
                pred, embed = model.forward(batch_graphs, batch_x, batch_e)

            # 将当前批次的真实标签和模型预测标签添加到之前初始化的列表中，并将它们从设备上移动到 CPU 上。
            y_true.append(batch[1].view(pred.shape).detach().cpu())
            y_pred.append(pred.detach().cpu())

    # 将所有真实标签和模型预测标签连接成一个大的张量，并转换为 NumPy 数组。
    y_true = torch.cat(y_true, dim = 0).numpy()
    y_pred = torch.cat(y_pred, dim = 0).numpy()

    # 将真实标签和模型预测标签存储在一个字典中
    input_dict = {"y_true": y_true, "y_pred": y_pred}

    # 使用评估器 evaluator 对模型的性能进行评估，并返回评估结果。
    return evaluator.eval(input_dict)


# 这个函数是训练和验证模型的主要流程。
# 它接受模型名称 MODEL_NAME、数据集名称 DATASET_NAME、数据集 dataset、训练参数 params、网络参数 net_params、目录信息 dirs、命令行参数 args 和输出目录 out_dir。
def train_val_pipeline(MODEL_NAME, DATASET_NAME, dataset, params, net_params, dirs, args, out_dir):
    # 记录训练过程中每个 epoch 所花费的时间。
    t0 = time.time()
    per_epoch_time = []
        
    # 针对 GCN 和 GAT 模型，如果网络参数中设置了 self_loop，就给图数据集添加自环。这是为了处理无向图中节点自身的信息。
    if MODEL_NAME in ['GCN', 'GAT']:
        if net_params['self_loop']:
            print("[!] Adding graph self-loops for GCN/GAT models (central node trick).")
            dataset._add_self_loops()
    
#    for i in range(len(trainset)):
#        print("num of nodes:", dataset.train[i][0].number_of_nodes())
    # 从目录信息中获取日志、模型检查点、结果文件和配置文件的根目录，并获取设备信息。
    root_log_dir, root_ckpt_dir, write_file_name, write_config_file = dirs
    device = net_params['device']

    # 将数据集、模型、参数和网络参数信息写入配置文件。
    # Write the network and optimization hyper-parameters in folder config/
    with open(write_config_file + '.txt', 'w') as f:
        f.write("""Dataset: {},\nModel: {}\n\nparams={}\n\nnet_params={}\n\n\nTotal Parameters: {}\n\n""".format("modelcular", MODEL_NAME, params, net_params, net_params['total_param']))

    # 根据日志目录创建一个 TensorBoard 日志写入器。
    log_dir = os.path.join(root_log_dir, "RUN_" + str(0))
    # 创建一个TensorboardX的SummaryWriter实例，用于记录训练过程中的各种指标。
    writer = SummaryWriter(log_dir=log_dir)

    # setting seeds
    # 设置随机种子以确保实验的可重复性。
    random.seed(params['seed'])
    np.random.seed(params['seed'])
    torch.manual_seed(params['seed'])
    if device.type == 'cuda':
        torch.cuda.manual_seed(params['seed'])
    
    #print("Training Graphs: ", len(trainset))
    #print("Validation Graphs: ", len(valset))
    #print("Test Graphs: ", len(biased_testset))
    # 打印出数据集中类别的数量。
    print("Number of Classes: ", net_params['n_classes'])

    # 获取数据集的索引分割，通常用于将数据集分为训练集、验证集和测试集。
    split_idx = dataset.get_idx_split()


    # 根据模型名称和网络参数创建模型。
    model = gnn_model(MODEL_NAME, net_params)
    #model = model.to(device)

    # 初始化优化器，这里使用 Adam 优化器。
    optimizer = optim.Adam(model.parameters(), lr=params['init_lr'])

    # 初始化训练和验证损失以及准确率的列表，用于记录每个 epoch 的指标。
    epoch_train_losses, epoch_val_losses = [], []
    epoch_train_accs, epoch_val_accs = [], [] 
    
    # batching exception for Diffpool
    # 如果模型是 DiffPool，则设置为 True，否则设置为 False。这是因为 DiffPool 模型的批次可能会有不同大小的图，因此在 DataLoader 中将最后一个批次丢弃。
    drop_last = True if MODEL_NAME == 'DiffPool' else False

    # 针对 RingGNN 和 3WLGNN 模型，使用特定的数据加载函数和评估函数，并创建对应的数据加载器。
    if MODEL_NAME in ['RingGNN', '3WLGNN']:
        # import train functions specific for WL-GNNs
        from train.train_superpixels_graph_classification import train_epoch_dense as train_epoch, evaluate_network_dense as evaluate_network

        train_loader = DataLoader(trainset, shuffle=True, collate_fn=dataset.collate_dense_gnn)
        val_loader = DataLoader(valset, shuffle=False, collate_fn=dataset.collate_dense_gnn)
        biased_test_loader = DataLoader(biased_testset, shuffle=False, collate_fn=dataset.collate_dense_gnn)
        unbiased_test_loader = DataLoader(unbiased_testset, shuffle=False, collate_fn=dataset.collate_dense_gnn)

    else:
        # import train functions for all other GCNs
        from train.train_superpixels_graph_classification import train_epoch_sparse as train_epoch, evaluate_network_sparse as evaluate_network, train_epoch_sparse_HSIC as train_epoch_HSIC

        ### automatic evaluator. takes dataset name as input

    # 对于其他模型，使用通用的数据加载函数和评估函数，并创建对应的数据加载器。
    train_loader = GraphDataLoader(dataset[split_idx["train"]], batch_size=net_params['batch_size'], shuffle=True, num_workers = 0, drop_last=drop_last)
    valid_loader = GraphDataLoader(dataset[split_idx["valid"]], batch_size=net_params['batch_size'], shuffle=True, num_workers = 0, drop_last=drop_last)
    test_loader = GraphDataLoader(dataset[split_idx["test"]], batch_size=net_params['batch_size'], shuffle=True, num_workers = 0, drop_last=drop_last)

    # 创建评估器，用于评估模型性能。
    evaluator = Evaluator(DATASET_NAME)
    Acc_evaluator = Evaluator("ogbg-ppa")
    F1_evaluator = Evaluator("ogbg-ppa")
    #F1_evaluator.eval_metric = "F1"

    # At any point you can hit Ctrl + C to break out of training early.
    # 初始化训练过程中的指标曲线和其他参数，包括验证集上的最佳性能、是否保存模型以及预训练的最大迭代次数。
    try:
        print("pretrain")
        valid_curve = []
        test_curve = []
        train_curve = []

        Acc_valid_curve = []
        Acc_test_curve = []
        Acc_train_curve = []

        F1_valid_curve = []
        F1_test_curve = []
        F1_train_curve = []

        best_val = 0
        # save_model变量用于控制是否保存模型。
        save_model=True
        #if 'pretrain' in params:
        #    pretrain_epoch = 0
        #else:
        pretrain_epoch = 100
        model = model.to(device)

        # 在预训练阶段，执行多个 epoch 的训练和验证，并记录训练过程中的损失和准确率。如果在验证集上获得了更好的性能，则保存模型。
        # 使用tqdm库创建一个进度条，用于显示预训练阶段的进度。range(pretrain_epoch)指定了预训练的总轮数。
        with tqdm(range(pretrain_epoch)) as t:
            # 开始一个循环，迭代预训练的每一轮。
            for epoch in t:
                # 设置进度条的描述，显示当前是第几轮。
                t.set_description('Epoch %d' % epoch)
                # 记录当前时间，用于计算每轮训练的耗时。
                start = time.time()

                # 根据模型名称，调用不同的训练函数来执行一轮训练。对于RingGNN和3WLGNN模型，使用特定的训练函数；对于其他模型，使用通用的训练函数。
                if MODEL_NAME in ['RingGNN', '3WLGNN']: # since different batch training function for dense GNNs
                    epoch_train_loss, epoch_train_acc, optimizer = train_epoch(model, optimizer, device, train_loader, epoch, params['batch_size'])
                else:   # for all other models common train function
                    epoch_train_loss, epoch_train_acc, optimizer = train_epoch(model, optimizer, device, train_loader, epoch, task_type=dataset.task_type)

                # 使用评估函数计算模型在训练集、验证集和测试集上的性能。
                train_perf = eval(model, device, train_loader, evaluator, Acc_evaluator, F1_evaluator)
                valid_perf = eval(model, device, valid_loader, evaluator, Acc_evaluator, F1_evaluator)
                test_perf = eval(model, device, test_loader, evaluator, Acc_evaluator, F1_evaluator)

                # 打印出模型在各个数据集上的性能指标。
                print({'Train': train_perf, 'Validation': valid_perf, 'Test': test_perf})

                # 将性能指标添加到相应的列表中，用于后续绘制性能曲线。
                train_curve.append(train_perf[dataset.eval_metric])
                valid_curve.append(valid_perf[dataset.eval_metric])
                test_curve.append(test_perf[dataset.eval_metric])
                
                # 如果save_model为True，则保存当前性能最佳的模型。
                if save_model:
                    ckpt_dir = os.path.join(root_ckpt_dir, "RUN_")
                    if not os.path.exists(ckpt_dir):
                        os.makedirs(ckpt_dir)
                    if valid_perf[dataset.eval_metric] > best_val:
                        torch.save(model.state_dict(), '{}.pkl'.format(ckpt_dir + "/epoch_best_mocular"))
                        best_val = valid_perf[dataset.eval_metric]

        # 如果预训练轮数大于0，则根据任务类型（分类或回归）选择最佳验证性能的轮数，并打印出最佳验证和测试性能。如果预训练轮数不大于0，则跳过这部分代码。
        if pretrain_epoch > 0: 
            if 'classification' in dataset.task_type:
                best_val_epoch = np.argmax(np.array(valid_curve))
                best_train = max(train_curve)
            else:
                best_val_epoch = np.argmin(np.array(valid_curve))
                best_train = min(train_curve)
            print('Finished training!')
            print('Best Auc validation score: {}, Test score: {}'.format(valid_curve[best_val_epoch], test_curve[best_val_epoch]))
            best_val_base = valid_curve[best_val_epoch]
            best_test_base = test_curve[best_val_epoch]
        else:
            pass
            #pretrain_dir = out_dir + "checkpoints/"+ params["pretrain"]
            #ckpt_dir = os.path.join(pretrain_dir, "RUN_")

        # 将模型设置为训练模式。
        model.train()
        # 再次将模型移动到指定的设备。
        model = model.to(device)
        #model.load_state_dict(torch.load('{}.pkl'.format(ckpt_dir + "/epoch_best_mocular")))

        # 重新创建优化器，这次包括了权重衰减。
        optimizer = optim.Adam(model.parameters(), lr=params['init_lr'], weight_decay=params['weight_decay'])
        # 创建一个学习率调度器，它会在验证损失不再改善时减少学习率。
        scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min',
                                                     factor=params['lr_reduce_factor'],
                                                     patience=params['lr_schedule_patience'],
                                                     verbose=True)


    
        print("our model start!")
        # 重新初始化性能指标列表。
        train_curve = []
        valid_curve = []
        test_curve = []
        # 创建一个新的进度条，用于显示训练的总轮数。
        with tqdm(range(100)) as t:
            # 开始一个新的循环，迭代训练的总轮数。
            for epoch in t:
                # 设置进度条的描述，显示当前是第几轮。
                t.set_description('Epoch %d' % epoch)
                # 记录当前时间，用于计算每轮训练的耗时。
                start = time.time()
                # 根据模型名称，调用不同的训练函数来执行一轮训练。对于RingGNN和3WLGNN模型，使用特定的训练函数；对于其他模型，使用通用的训练函数。
                if MODEL_NAME in ['RingGNN', '3WLGNN']: # since different batch training function for dense GNNs
                    epoch_train_loss, epoch_train_acc, optimizer = train_epoch(model, optimizer, device, train_loader, epoch, params['batch_size'])
                else:   # for all other models common train function
                    epoch_train_loss, epoch_train_acc, optimizer = train_epoch_HSIC(model, optimizer, net_params['hidden_dim'], net_params["assign_num"], device, train_loader, epoch, params["lrbl"], params["lambdap"], args.lambda_decay_rate, args.lambda_decay_epoch, args.min_lambda_times, task_type=dataset.task_type)

                # 使用评估函数计算模型在训练集、验证集和测试集上的性能。
                train_perf = eval(model, device, train_loader, evaluator, Acc_evaluator, F1_evaluator)
                valid_perf = eval(model, device, valid_loader, evaluator, Acc_evaluator, F1_evaluator)
                test_perf = eval(model, device, test_loader, evaluator, Acc_evaluator, F1_evaluator)

                # 打印出模型在各个数据集上的性能指标。
                print({'Train': train_perf, 'Validation': valid_perf, 'Test': test_perf})
                # 将性能指标添加到相应的列表中，用于后续绘制性能曲线。
                train_curve.append(train_perf[dataset.eval_metric])
                valid_curve.append(valid_perf[dataset.eval_metric])
                test_curve.append(test_perf[dataset.eval_metric])
                
        # 根据任务类型选择最佳验证性能的轮数，并打印出最佳验证和测试性能。
        if 'classification' in dataset.task_type:
            best_val_epoch = np.argmax(np.array(valid_curve))
            best_train = max(train_curve)
        else:
            best_val_epoch = np.argmin(np.array(valid_curve))
            best_train = min(train_curve)


        print('Finished training!')
        print('Best Auc validation score: {}, Test score: {}'.format(valid_curve[best_val_epoch], test_curve[best_val_epoch]))

        best_val_our = valid_curve[best_val_epoch]
        best_test_our = test_curve[best_val_epoch]
        

    except KeyboardInterrupt:
        print('-' * 89)
        print('Exiting from training early because of KeyboardInterrupt')
    # 将最终结果写入结果文件。
    with open(write_file_name + '.txt', 'w') as f:
        f.write("""Dataset: {},\nModel: {}\n\nparams={}\n\nnet_params={}\n\n{}\n\nTotal Parameters: {}\n\n
    FINAL RESULTS\nbest_validation_base: {:.4f}\n\nbest_test_base: {:.4f}\n\nbest_validation_our: {:.4f}\n\nbest_test_base: {:.4f}\n\n
    Convergence Time (Epochs): {:.4f}\nTotal Time Taken: {:.4f} hrs\nAverage Time Per Epoch: {:.4f} s\n\n\n"""\
          .format(DATASET_NAME, MODEL_NAME, params, net_params, model, net_params['total_param'],
                   best_val_base*100,  best_test_base*100, best_val_our*100, best_test_our*100, epoch, (time.time()-t0)/3600, np.mean(per_epoch_time)))
 

def main():    
    """
        USER CONTROLS
    """
    
    # 创建了一个参数解析器对象。
    parser = argparse.ArgumentParser()
    # 这些代码行定义了程序所需要的命令行参数，每个参数都有相应的帮助文本。
    parser.add_argument('--config', help="Please give a config.json file with training/model/data/param details")
    parser.add_argument('--gpu_id', help="Please give a value for gpu id")
    parser.add_argument('--model', help="Please give a value for model name")
    parser.add_argument('--dataset', help="Please give a value for dataset name")
    parser.add_argument('--train_str', help="Please give a value for dataset name")
    parser.add_argument('--out_dir', help="Please give a value for out_dir")
    parser.add_argument('--seed', help="Please give a value for seed")
    parser.add_argument('--epochs', help="Please give a value for epochs")
    parser.add_argument('--batch_size', help="Please give a value for batch_size")
    parser.add_argument('--init_lr', help="Please give a value for init_lr")
    parser.add_argument('--lr_reduce_factor', help="Please give a value for lr_reduce_factor")
    parser.add_argument('--lr_schedule_patience', help="Please give a value for lr_schedule_patience")
    parser.add_argument('--min_lr', help="Please give a value for min_lr")
    parser.add_argument('--weight_decay', help="Please give a value for weight_decay")
    parser.add_argument('--print_epoch_interval', help="Please give a value for print_epoch_interval")    
    parser.add_argument('--L', help="Please give a value for L")
    parser.add_argument('--hidden_dim', help="Please give a value for hidden_dim")
    parser.add_argument('--out_dim', help="Please give a value for out_dim")
    parser.add_argument('--residual', help="Please give a value for residual")
    parser.add_argument('--edge_feat', help="Please give a value for edge_feat")
    parser.add_argument('--readout', help="Please give a value for readout")
    parser.add_argument('--kernel', help="Please give a value for kernel")
    parser.add_argument('--n_heads', help="Please give a value for n_heads")
    parser.add_argument('--gated', help="Please give a value for gated")
    parser.add_argument('--in_feat_dropout', help="Please give a value for in_feat_dropout")
    parser.add_argument('--dropout', help="Please give a value for dropout")
    parser.add_argument('--layer_norm', help="Please give a value for layer_norm")
    parser.add_argument('--batch_norm', help="Please give a value for batch_norm")
    parser.add_argument('--sage_aggregator', help="Please give a value for sage_aggregator")
    parser.add_argument('--data_mode', help="Please give a value for data_mode")
    parser.add_argument('--num_pool', help="Please give a value for num_pool")
    parser.add_argument('--gnn_per_block', help="Please give a value for gnn_per_block")
    parser.add_argument('--embedding_dim', help="Please give a value for embedding_dim")
    parser.add_argument('--pool_ratio', help="Please give a value for pool_ratio")
    parser.add_argument('--linkpred', help="Please give a value for linkpred")
    parser.add_argument('--cat', help="Please give a value for cat")
    parser.add_argument('--self_loop', help="Please give a value for self_loop")
    parser.add_argument('--max_time', help="Please give a value for max_time")
    parser.add_argument('--lambda_decay_rate', type=float, default=1, help = 'ratio of epoch for lambda to decay')
    parser.add_argument('--lambda_decay_epoch', type=int, default=5, help = 'number of epoch for lambda to decay')
    parser.add_argument('--min_lambda_times', type=float, default=0.01, help = 'number of global table levels')
    parser.add_argument('--lambdap', type = float, default = 2.0, help = 'weight decay for weight1 ')
    parser.add_argument('--lrbl', type = float, default = 1.0, help = 'weight decay for weight1 ')
    parser.add_argument('--pretrain', help = 'weight decay for weight1 ')
    # 这行代码解析命令行参数，并将其存储在 args 变量中。
    args = parser.parse_args()
    # 这行代码打开并读取指定文件中的 JSON 配置信息，并将其存储在 config 变量中。
    with open(args.config) as f:
        config = json.load(f)
        
    # device
    # 这段代码根据命令行参数中指定的 GPU ID，设置 GPU 使用信息，并调用 gpu_setup 函数设置使用的 GPU 设备。
    if args.gpu_id is not None:
        config['gpu']['id'] = int(args.gpu_id)
        config['gpu']['use'] = True
    device = gpu_setup(config['gpu']['use'], config['gpu']['id'])
    # model, dataset, out_dir
    # 这段代码根据命令行参数或配置文件中指定的模型名称和数据集名称，设置 MODEL_NAME 和 DATASET_NAME 变量。
    if args.model is not None:
        MODEL_NAME = args.model
    else:
        MODEL_NAME = config['model']
    if args.dataset is not None:
        DATASET_NAME = args.dataset
    else:
        DATASET_NAME = config['dataset']

    # 这行代码使用指定名称的数据集创建 DglGraphPropPredDataset 对象。
    dataset = DglGraphPropPredDataset(name = args.dataset)

    # 这行代码打印数据集中第一个样本的信息。
    print("data 0 ", dataset[0])

    #dataset = LoadData(DATASET_NAME, args.train_str, '1_biased', '0.1_unbiased')
    # 这段代码根据命令行参数或配置文件中指定的输出目录，设置 out_dir 变量。
    if args.out_dir is not None:
        out_dir = args.out_dir
    else:
        out_dir = config['out_dir']
    # parameters
    # 这行代码将配置文件中的参数部分存储在 params 变量中。
    params = config['params']
    # 接下来一系列的 if 语句根据命令行参数设置相应的参数值。
    if args.seed is not None:
        params['seed'] = int(args.seed)
    if args.epochs is not None:
        params['epochs'] = int(args.epochs)
    if args.batch_size is not None:
        params['batch_size'] = int(args.batch_size)
    if args.init_lr is not None:
        params['init_lr'] = float(args.init_lr)
    if args.lr_reduce_factor is not None:
        params['lr_reduce_factor'] = float(args.lr_reduce_factor)
    if args.lr_schedule_patience is not None:
        params['lr_schedule_patience'] = int(args.lr_schedule_patience)
    if args.min_lr is not None:
        params['min_lr'] = float(args.min_lr)
    if args.weight_decay is not None:
        params['weight_decay'] = float(args.weight_decay)
    if args.print_epoch_interval is not None:
        params['print_epoch_interval'] = int(args.print_epoch_interval)
    if args.max_time is not None:
        params['max_time'] = float(args.max_time)
    if args.lrbl is not None:
        params['lrbl'] = float(args.lrbl)
    if args.lambdap is not None:
        params['lambdap'] = float(args.lambdap)
    if args.pretrain is not None:
        params['pretrain'] = args.pretrain

    # network parameters
    # 这行代码将配置文件中的网络参数部分存储在 net_params 变量中。
    net_params = config['net_params']

    # 接下来一系列的 if 语句根据命令行参数设置相应的网络参数值。
    net_params['in_dim'] = dataset[0][0].ndata['feat'].size(1)
    net_params['device'] = device
    net_params['gpu_id'] = config['gpu']['id']
    net_params['batch_size'] = params['batch_size']
    if args.L is not None:
        net_params['L'] = int(args.L)
    if args.hidden_dim is not None:
        net_params['hidden_dim'] = int(args.hidden_dim)
    if args.out_dim is not None:
        net_params['out_dim'] = int(args.out_dim)   
    if args.residual is not None:
        net_params['residual'] = True if args.residual=='True' else False
    if args.edge_feat is not None:
        net_params['edge_feat'] = True if args.edge_feat=='True' else False
    if args.readout is not None:
        net_params['readout'] = args.readout
    if args.kernel is not None:
        net_params['kernel'] = int(args.kernel)
    if args.n_heads is not None:
        net_params['n_heads'] = int(args.n_heads)
    if args.gated is not None:
        net_params['gated'] = True if args.gated=='True' else False
    if args.in_feat_dropout is not None:
        net_params['in_feat_dropout'] = float(args.in_feat_dropout)
    if args.dropout is not None:
        net_params['dropout'] = float(args.dropout)
    if args.layer_norm is not None:
        net_params['layer_norm'] = True if args.layer_norm=='True' else False
    if args.batch_norm is not None:
        net_params['batch_norm'] = True if args.batch_norm=='True' else False
    if args.sage_aggregator is not None:
        net_params['sage_aggregator'] = args.sage_aggregator
    if args.data_mode is not None:
        net_params['data_mode'] = args.data_mode
    if args.num_pool is not None:
        net_params['num_pool'] = int(args.num_pool)
    if args.gnn_per_block is not None:
        net_params['gnn_per_block'] = int(args.gnn_per_block)
    if args.embedding_dim is not None:
        net_params['embedding_dim'] = int(args.embedding_dim)
    if args.pool_ratio is not None:
        net_params['pool_ratio'] = float(args.pool_ratio)


    if args.linkpred is not None:
        net_params['linkpred'] = True if args.linkpred=='True' else False
    if args.cat is not None:
        net_params['cat'] = True if args.cat=='True' else False
    if args.self_loop is not None:
        net_params['self_loop'] = True if args.self_loop=='True' else False
        
    # Superpixels
    #net_params['in_dim'] = dataset.train[0][0].ndata['feat'][0].size(0)
    #net_params['in_dim_edge'] = dataset.train[0][0].edata['feat'][0].size(0)
    net_params['n_classes'] = dataset.num_tasks

    if MODEL_NAME == 'DiffPool':
        # calculate assignment dimension: pool_ratio * largest graph's maximum
        # number of nodes  in the dataset
        max_num_nodes_train = max([dataset[i][0].number_of_nodes() for i in range(len(dataset))])
        max_num_node = max_num_nodes_train
        #net_params['assign_dim'] = int(max_num_node * net_params['pool_ratio']) * net_params['batch_size']
        #net_params["assign_num"] = int(max_num_node * net_params['pool_ratio'])

        net_params["assign_num"] = 7
        net_params['assign_dim'] = net_params["assign_num"] * net_params['batch_size']
    else:
        net_params["assign_num"] = 0
        net_params["pool_ratio"] = 0
        
        
    if MODEL_NAME == 'RingGNN':
        num_nodes_train = [dataset.train[i][0].number_of_nodes() for i in range(len(dataset.train))]
        num_nodes_test = [dataset.test[i][0].number_of_nodes() for i in range(len(dataset.test))]
        num_nodes = num_nodes_train + num_nodes_test
        net_params['avg_node_num'] = int(np.ceil(np.mean(num_nodes)))
        
    # 这段代码生成日志目录、检查点目录、结果文件名和配置文件名，并创建相应的目录。
    root_log_dir = out_dir + 'logs/' + MODEL_NAME + "_" + DATASET_NAME + "_GPU" + str(config['gpu']['id']) + "_" + time.strftime('%Hh%Mm%Ss_on_%b_%d_%Y') +'_'+str(args.train_str)+"_"+str(params['seed'])
    root_ckpt_dir = out_dir + 'checkpoints/' + MODEL_NAME + "_" + DATASET_NAME + "_GPU" + str(config['gpu']['id']) + "_" + time.strftime('%Hh%Mm%Ss_on_%b_%d_%Y')  +'_'+str(args.train_str)+"_"+str(params['seed'])
    write_file_name = out_dir + 'results/result_' + MODEL_NAME + "_" + DATASET_NAME + "_GPU" + str(config['gpu']['id']) + "_" + time.strftime('%Hh%Mm%Ss_on_%b_%d_%Y')  +'_'+str(args.train_str) + '_assign_ratio_'+str(net_params['pool_ratio']) +"_"+str(params['seed'])
    write_config_file = out_dir + 'configs/config_' + MODEL_NAME + "_" + DATASET_NAME + "_GPU" + str(config['gpu']['id']) + "_" + time.strftime('%Hh%Mm%Ss_on_%b_%d_%Y') +'_'+str(args.train_str) +"_"+str(params['seed'])
    dirs = root_log_dir, root_ckpt_dir, write_file_name, write_config_file

    if not os.path.exists(out_dir + 'results'):
        os.makedirs(out_dir + 'results')
        
    if not os.path.exists(out_dir + 'configs'):
        os.makedirs(out_dir + 'configs')

    # 这行代码调用 view_model_param 函数获取模型参数数量，并将结果存储在 net_params['total_param'] 中。
    net_params['total_param'] = view_model_param(MODEL_NAME, net_params)
    # 这行代码调用 train_val_pipeline 函数开始模型的训练和验证过程，传递了模型名称、数据集名称、数据集对象、参数、网络参数、目录信息等参数。
    train_val_pipeline(MODEL_NAME, args.dataset, dataset, params, net_params, dirs, args, out_dir)

    
    
    
main()    
















